use colored::Colorize;
use log::debug;
use std::cell::{RefCell};
use std::fmt;
use std::fmt::Debug;
use std::rc::{Rc, Weak};

use pest::Parser;
use pest::error::{Error, LineColLocation};
use pest::iterators::{Pair, Pairs};
use pest_derive::Parser;

#[derive(Parser)]
#[grammar = "parser2.pest"]
pub struct ExpressionParser;

#[derive(Debug, PartialEq, Clone)]
pub enum Token {
    Root,
    KeyWord(String),
    Plus(String),
    Minus(String),
    Mul(String),
    Div(String),
    Mod(String),
    Assign(String),
    Eq(String),
    Neq(String),
    Lt(String),
    Gt(String),
    Le(String),
    Ge(String),
    Not(String),
    And(String),
    Or(String),
    LParen(String),
    RParen(String),
    LBrace(String),
    RBrace(String),
    LBrackt(String),
    RBrackt(String),
    Comma(String),
    Semicolon(String),

    Ident(String),
    CompUnit,
    Decl,
    // ConstDecl,
    BType(String),
    // ConstDef,
    // ConstInitVal,
    // VarDecl,
    // VarDef,
    // InitVal,
    FuncDef,
    FuncType(String),
    FuncFParams,
    FuncFParam,
    Block,
    // BlockItem,
    IfStmt,
    WhileStmt,
    Stmt,
    Exp,
    // Cond,
    // LVal,
    PrimaryExp,
    Number(String),
    UnaryExp,
    UnaryOp,
    // FuncRParams,
    // MulExp,
    // AddExp,
    // RelExp,
    // EqExp,
    // LAndExp,
    // LOrExp,
    ConstExp,
}

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Token::Root => write!(f, "Root"),
            Token::KeyWord(s) => write!(f, "KeyWord {}", s.yellow()),
            Token::Plus(s) => write!(f, "Plus {}", s.yellow()),
            Token::Minus(s) => write!(f, "Minus {}", s.yellow()),
            Token::Mul(s) => write!(f, "Mul {}", s.yellow()),
            Token::Div(s) => write!(f, "Div {}", s.yellow()),
            Token::Mod(s) => write!(f, "Mod {}", s.yellow()),
            Token::Assign(s) => write!(f, "Assign {}", s.yellow()),
            Token::Eq(s) => write!(f, "Eq {}", s.yellow()),
            Token::Neq(s) => write!(f, "Neq {}", s.yellow()),
            Token::Lt(s) => write!(f, "Lt {}", s.yellow()),
            Token::Gt(s) => write!(f, "Gt {}", s.yellow()),
            Token::Le(s) => write!(f, "Le {}", s.yellow()),
            Token::Ge(s) => write!(f, "Ge {}", s.yellow()),
            Token::Not(s) => write!(f, "Not {}", s.yellow()),
            Token::And(s) => write!(f, "And {}", s.yellow()),
            Token::Or(s) => write!(f, "Or {}", s.yellow()),
            Token::LParen(s) => write!(f, "LParen {}", s.yellow()),
            Token::RParen(s) => write!(f, "RParen {}", s.yellow()),
            Token::LBrace(s) => write!(f, "LBrace {}", s.yellow()),
            Token::RBrace(s) => write!(f, "RBrace {}", s.yellow()),
            Token::LBrackt(s) => write!(f, "LBrackt {}", s.yellow()),
            Token::RBrackt(s) => write!(f, "RBrackt {}", s.yellow()),
            Token::Comma(s) => write!(f, "Comma {}", s.yellow()),
            Token::Semicolon(s) => write!(f, "Semicolon {}", s.yellow()),
            Token::Ident(s) => write!(f, "Ident {}", s.yellow()),
            Token::CompUnit => write!(f, "CompUnit"),
            Token::Decl => write!(f, "Decl"),
            Token::BType(s) => write!(f, "BType {}", s.yellow()),
            Token::FuncDef => write!(f, "FuncDef"),
            Token::FuncType(s) => write!(f, "FuncType {}", s.yellow()),
            Token::FuncFParams => write!(f, "FuncFParams"),
            Token::FuncFParam => write!(f, "FuncFParam"),
            Token::Block => write!(f, "Block"),
            Token::IfStmt => write!(f, "IfStmt"),
            Token::WhileStmt => write!(f, "WhileStmt"),
            Token::Stmt => write!(f, "Stmt"),
            Token::Exp => write!(f, "Exp"),
            Token::PrimaryExp => write!(f, "PrimaryExp"),
            Token::Number(s) => write!(f, "Number {}", s.yellow()),
            Token::UnaryExp => write!(f, "UnaryExp"),
            Token::UnaryOp => write!(f, "UnaryOp"),
            Token::ConstExp => write!(f, "ConstExp"),
        }
    }
}

pub struct Node {
    pub token: Token,
    pub parent: Option<Weak<RefCell<Node>>>,
    pub children: Vec<Option<Rc<RefCell<Node>>>>,
}

impl Node {
    pub fn new() -> Self {
        Node {
            token: Token::Root,
            parent: None,
            children: vec![],
        }
    }

    pub fn insert(&mut self, node: Rc<RefCell<Node>>) {
        self.children.push(Some(node));
    }

    pub fn has_parent(&self, token: Token, max_level: usize) -> bool {
        let mut current = self.parent.as_ref().and_then(|weak| weak.upgrade());
        let mut level = 0;
        while let Some(parent_rc) = current {
            let parent = parent_rc.borrow();
            if parent.token == token {
                return true;
            }
            current = parent.parent.as_ref().and_then(|weak| weak.upgrade());
            level += 1;
            if level >= max_level {
                return false;
            }
        }
        false
    }

    pub fn print_ast(
        &self,
        f: &mut fmt::Formatter,
        depth: usize,
    ) -> fmt::Result {
        let indent = "  ".repeat(depth);
        write!(f, "{}{}\n", indent, self.token)?;
        for child in self.children.iter() {
            if let Some(node) = child {
                node.borrow().print_ast(f, depth + 1)?;
            }
        }
        Ok(())
    }
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "===== Abstract syntax tree =====\n")?;
        self.print_ast(f, 0)?;
        Ok(())
    }
}

#[allow(dead_code)]
fn print_error(input: &str, error: Error<Rule>) {
    let (line, col) = match error.line_col {
        LineColLocation::Pos((line, col)) => (line, col),
        LineColLocation::Span((line, col), _) => (line, col),
    };
    let lines: Vec<&str> = input.lines().collect();
    if line <= lines.len() {
        let problematic_line = lines[line - 1];
        let char = problematic_line.chars().nth(col - 1).unwrap();
        eprintln!(
            "Error type A at Line {}: Mysterious character \"{}\".",
            line, char
        );
    } else {
        eprintln!("Invalid line={}/{}", line, lines.len());
    }
    // std::process::exit(1);
}

fn visit_sub_node(
    ast: &mut Rc<RefCell<Node>>,
    pair: Pair<Rule>,
    depth: usize,
    token: Option<Token>,
) {
    if let Some(t) = token {
        let mut sub_ast = Rc::new(RefCell::new(Node {
            token: t,
            parent: Some(Rc::downgrade(ast)),
            children: vec![],
        }));
        // Add sub nodes.
        for sub_pair in pair.into_inner() {
            visit(&mut sub_ast, sub_pair, depth + 1);
        }
        // Record as current ast node.
        ast.borrow_mut().insert(sub_ast);
    } else {
        // Skip current token and parse subtree.
        for sub_pair in pair.into_inner() {
            visit(ast, sub_pair, depth + 1);
        }
    }
}

macro_rules! token {
    ($token: expr) => {
        Some($token)
    };

    ($token: expr, $pair: expr) => {
        Some($token($pair.as_str().to_string()))
    };
}

// Recursively print all inner rules
fn visit(ast: &mut Rc<RefCell<Node>>, pair: Pair<Rule>, depth: usize) {
    let indent = "  ".repeat(depth);
    let (line, _) = pair.line_col();
    debug!(
        "{} rule={:?} str={} line={}.",
        indent,
        pair.as_rule(),
        pair.as_str(),
        line
    );
    let token = match pair.as_rule() {
        Rule::IF => token!(Token::KeyWord, pair),
        Rule::ELSE => token!(Token::KeyWord, pair),
        Rule::WHILE => token!(Token::KeyWord, pair),
        Rule::BREAK => token!(Token::KeyWord, pair),
        Rule::CONTINUE => token!(Token::KeyWord, pair),
        Rule::RETURN => token!(Token::KeyWord, pair),
        Rule::PLUS => token!(Token::Plus, pair),
        Rule::MINUS => token!(Token::Minus, pair),
        Rule::MUL => token!(Token::Mul, pair),
        Rule::DIV => token!(Token::Div, pair),
        Rule::MOD => token!(Token::Mod, pair),
        Rule::ASSIGN => token!(Token::Assign, pair),
        Rule::EQ => token!(Token::Eq, pair),
        Rule::NEQ => token!(Token::Neq, pair),
        Rule::LT => token!(Token::Lt, pair),
        Rule::GT => token!(Token::Gt, pair),
        Rule::LE => token!(Token::Le, pair),
        Rule::GE => token!(Token::Ge, pair),
        Rule::NOT => token!(Token::Not, pair),
        Rule::AND => token!(Token::And, pair),
        Rule::OR => token!(Token::Or, pair),
        Rule::L_PAREN => token!(Token::LParen, pair),
        Rule::R_PAREN => token!(Token::RParen, pair),
        Rule::L_BRACE => token!(Token::LBrace, pair),
        Rule::R_BRACE => token!(Token::RBrace, pair),
        Rule::L_BRACKT => token!(Token::LBrackt, pair),
        Rule::R_BRACKT => token!(Token::RBrackt, pair),
        Rule::COMMA => token!(Token::Comma, pair),
        Rule::SEMICOLON => token!(Token::Semicolon, pair),
        Rule::Ident => token!(Token::Ident, pair),
        Rule::CompUnit => token!(Token::CompUnit),
        Rule::Decl => token!(Token::Decl),
        Rule::BType => token!(Token::BType, pair),
        Rule::FuncDef => token!(Token::FuncDef),
        Rule::FuncType => token!(Token::FuncType, pair),
        Rule::FuncFParams => token!(Token::FuncFParams),
        Rule::FuncFParam => token!(Token::FuncFParam),
        Rule::Block => token!(Token::Block),
        Rule::IfStmt => token!(Token::IfStmt),
        Rule::WhileStmt => token!(Token::WhileStmt),
        Rule::Stmt => token!(Token::Stmt),
        Rule::Exp => token!(Token::Exp),
        // Rule::PrimaryExp => token!(Token::PrimaryExp),
        Rule::Number => token!(Token::Number, pair),
        // Rule::UnaryExp => token!(Token::UnaryExp),
        Rule::UnaryOp => token!(Token::UnaryOp),
        Rule::ConstExp => token!(Token::ConstExp),
        _ => None,
    };
    visit_sub_node(ast, pair, depth, token);
}

pub fn tokenize(input: &str) -> Rc<RefCell<Node>> {
    // DEBUG
    // let pairs = ExpressionParser::parse(Rule::sy, input)
    //     .unwrap_or_else(|e| panic!("{}", e));
    let res = ExpressionParser::parse(Rule::sy, input);
    let mut ast = Rc::new(RefCell::new(Node::new()));
    if res.is_ok() {
        let pairs = res.unwrap();
        for pair in pairs {
            visit(&mut ast, pair, 0);
            // let (line, _) = pair.line_col();
            // let rule = pair.as_rule();
            // let content = pair.as_str();
            // let content_output = match rule {
            //     Rule::INTEGER_CONST => parse_num(content)
            //         .unwrap_or_else(|e| {
            //             panic!("Failed to parse={} err={}", content, e)
            //         })
            //         .to_string(),
            //     _ => content.to_string(),
            // };
            // eprintln!("{:?} {} at Line {}.", rule, content_output, line);
        }
    } else {
        print_error(input, res.err().unwrap());
    }
    ast
}

pub fn parse_with_recovery(input: &str) -> Vec<Result<Pairs<Rule>, String>> {
    let mut results = Vec::new();
    let mut remaining_input = input;

    while !remaining_input.is_empty() {
        match ExpressionParser::parse(Rule::sy, remaining_input) {
            Ok(pairs) => {
                let consumed = pairs.as_str();
                results.push(Ok(pairs));
                remaining_input = &remaining_input[consumed.len()..];
            }
            Err(error) => {
                let (line, col) = match error.line_col {
                    LineColLocation::Pos((line, col)) => (line, col),
                    LineColLocation::Span((line, col), _) => (line, col),
                };
                eprintln!(
                    "Error at line {}, column {} {}",
                    line,
                    col,
                    error.to_string()
                );

                // Try to recover by skipping to next semicolon or newline
                if let Some(next_semi) = remaining_input.find(';') {
                    remaining_input = &remaining_input[next_semi + 1..];
                } else if let Some(next_newline) = remaining_input.find('\n') {
                    remaining_input = &remaining_input[next_newline + 1..];
                } else {
                    break; // Cannot recover further
                }
            }
        }
    }
    results
}
